from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr
from matplotlib.widgets import Button
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from scipy.optimize import curve_fit
from scipy.fft import fft2,ifft2
from scipy.fft import fftn, ifftn
from scipy.interpolate import Rbf
import scipy.signal
from scipy.ndimage import gaussian_filter
from ripser import ripser
from persim import plot_diagrams
from scipy.spatial.distance import pdist, squareform

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

# Parameters
image_size = 41
num_nodes = 38
sigma = 1
peak_value = 10
sample_start_x, sample_start_y = 65, 48
sample_size = 41
step = 2
sample_range = 100

# imgs = generate_maze_pool(num_mazes=1000, width=image_size, height=image_size, weight=0.2)
# imgs = np.array(imgs)
# print(imgs.shape)

# for i in range(imgs.shape[0]):
#     progress_bar(i, imgs.shape[0])
#     img = gaussian_filter(imgs[i], sigma=sigma)
#     imgs[i] = img.copy()

# # 生成10000张空白图像
# imgs = np.zeros((image_size*image_size, image_size, image_size))

# # 在每一个位置上生成一个 peak，模拟 pace field
# for i in range(image_size*image_size):
#     progress_bar(i, image_size*image_size)
#     x = i-(i//image_size)*image_size
#     y = i//image_size
#     img = np.zeros((image_size, image_size))
#     img[x, y] = peak_value
#     img = gaussian_filter(img, sigma=sigma)
#     imgs[i] = img

n_imgs = 4000
imgs = np.zeros((n_imgs, image_size, image_size))

# 在每一个位置上生成一个 peak，模拟 pace field
sigma = 0.1
for i in range(n_imgs):
    progress_bar(i, n_imgs)
    img = np.zeros((image_size, image_size))
    # x = random.randint(0, image_size-1)
    # y = random.randint(0, image_size-1)
    # img[x, y] = peak_value
    xs = [random.randint(0, image_size-1) for _ in range(random.randint(1, 50))]
    ys = [random.randint(0, image_size-1) for _ in range(random.randint(1, 50))]
    for x, y in zip(xs, ys):
        img[x, y] = peak_value
    img = gaussian_filter(img, sigma=sigma)
    imgs[i] = img

# # 在每一个位置上生成随机噪音
# sigma = 0.01
# for i in range(n_imgs):
#     progress_bar(i, n_imgs)
#     img = np.random.rand(image_size, image_size)
#     img = gaussian_filter(img, sigma=sigma)
#     imgs[i] = img

# 随机选择一张图像，显示出来
idx = random.randint(0, imgs.shape[0]-1)
plt.imshow(imgs[idx], cmap='viridis', interpolation='nearest')
plt.show()
    
# 将 imgs 的后两维拉平
imgs_reshaped = imgs.reshape((imgs.shape[0], imgs.shape[1]*imgs.shape[2]))

pca = PCA()
pca.fit(imgs_reshaped)
imgs_reshaped_pca = pca.transform(imgs_reshaped)

imgs_reshaped_pca_components = pca.components_

# 将 ridge_images_pca_grouped 绘制到3D视图中
fig = plt.figure()
ax1 = fig.add_subplot(111)

view_dim0, view_dim1, view_dim2 = 0, 1, 2

def button_clicked1(event):

    global view_dim0, view_dim1, view_dim2
    if view_dim2 <= imgs_reshaped_pca_components.shape[0]-2:
        view_dim0 += 1
    # 在这里清空原有的图像
    ax1.clear()
    
    # 显示 ridge_pca_components[view_dim0-2] 的 21*21 图像
    ax1.imshow(imgs_reshaped_pca_components[view_dim0].reshape(image_size,image_size))
    
    ax1.set_title('view_dim0: ' + str(view_dim0))
    # 刷新图像
    plt.draw()


def button_clicked2(event):
    global view_dim0, view_dim1, view_dim2
    if view_dim0 > 0:
        view_dim0 -= 1
    # 在这里清空原有的图像
    ax1.clear()
    
    # 显示 ridge_pca_components[view_dim0-2] 的 21*21 图像
    ax1.imshow(imgs_reshaped_pca_components[view_dim0].reshape(image_size,image_size))
    
    ax1.set_title('view_dim0: ' + str(view_dim0))
    # 刷新图像
    plt.draw()

# 创建按钮的位置和大小
button_ax1 = fig.add_axes([0.6, 0.1, 0.1, 0.1])
button_ax2 = fig.add_axes([0.3, 0.1, 0.1, 0.1])
# 创建按钮对象，并设置按钮的文本和位置
button1 = Button(button_ax1, 'dim_forward')
button2 = Button(button_ax2, 'dim_backward')
# 绑定按钮的点击事件和响应函数
button1.on_clicked(button_clicked1)
button2.on_clicked(button_clicked2)

# 显示 ridge_pca_components[view_dim0-2] 的 21*21 图像
ax1.imshow(imgs_reshaped_pca_components[view_dim0].reshape(image_size,image_size))

ax1.set_title('view_dim0: ' + str(view_dim0))
# 刷新图像
plt.show()

# 绘制 PCA 的 variance ratio
fig = plt.figure()
ax = fig.add_subplot(111)
ax.bar(range(len(pca.explained_variance_ratio_)), pca.explained_variance_ratio_)
ax.set_title('explained variance ratio')
plt.show()